import cv2
from pprint import pprint

import numpy as np

# 计算算术均值
def calc_arithmetic_mean(data, mask):
    return np.sum(data * mask) / np.sum(mask)

# 计算几何均值
def calc_geometric_mean(data, mask):
    size = len(data)
    res = 1
    for i in range(size):
        for j in range(size):
            res *= data[i, j] ** (1 / size ** 2)
    return int(res + 0.5)

# 计算谐波均值
def calc_harmonic_mean(data, mask):
    size = len(data)
    res = 0
    for i in range(size):
        for j in range(size):
            if data[i, j] == 0:
                continue
            res += 1 / data[i, j]
    return int(size ** 2 / res)

# 计算反谐波均值
def calc_contra_harmonic_mean(data, mask):
    size = len(data)
    res1 = np.float_(0)
    res2 = np.float_(0)
    for i in range(size):
        for j in range(size):
            if data[i, j] == 0:
                continue
            res1 += data[i, j] ** (mask + 1)
            res2 += data[i, j] ** mask
    if res1 == np.inf or res2 == np.inf:
        return 0
    return int(res1 / res2)

# 中值滤波
def calc_median(data, mask):
    size = len(data)
    tmp = []
    for i in range(size):
        for j in range(size):
            tmp.append(data[i, j])
    tmp.sort()
    size = len(tmp)
    if size % 2 == 0:
        return tmp[size // 2] / 2 + tmp[size // 2 - 1] / 2
    else:
        return tmp[size // 2]

# 滤波主体函数，通过调用不同的计算方法来实现不同的均值计算
def mean_filter(ori, mask, method, size):
    H, W, C = np.shape(ori)
    nH, nW = H - size + 1, W - size + 1
    ori = ori[:, :, 0]
    new = np.zeros((nH, nW, 3), dtype=np.uint8)
    for i in range(nH):
        for j in range(nW):
            new[i, j, :] = method(ori[i:i + size, j:j + size], mask)
    return new


# 计算mse
def mse_loss(ori, img):
    ori = ori[:, :, 0]
    img = img[:, :, 0]
    H1, W1 = np.shape(ori)
    H2, W2 = np.shape(img)
    pad = int((H1 - H2) / 2)
    # pprint(pad)
    res = np.float64(0)
    for i in range(H2):
        for j in range(W2):
            res += (ori[i + pad, j + pad] - img[i, j]) ** 2
    res = res / (H2 ** 2)
    return res


if __name__ == '__main__':
    # 读取原图像并转换数据类型为浮点型
    ori = cv2.imread('./test/test2_1_org.jpg')
    ori = ori.astype('float64')

    noise1 = cv2.imread('./test/test2_1.jpg')
    noise1 = noise1.astype('float64')

    # 使用3*3的mask的均值滤波
    mask = np.ones((3, 3))
    new = mean_filter(noise1, mask, calc_arithmetic_mean, 3)
    cv2.imwrite('./res/exp2/res_test2_1_ones_3.jpg', new)
    pprint("使用3*3的mask的均值滤波： " + str(mse_loss(ori, new)))

    # 使用5*5的mask的均值滤波
    mask = np.ones((5, 5))
    new = mean_filter(noise1, mask, calc_arithmetic_mean, 5)
    cv2.imwrite('./res/exp2/res_test2_1_ones_5.jpg', new)
    pprint("使用5*5的mask的均值滤波： " + str(mse_loss(ori, new)))

    # 使用3*3的加权均值滤波
    mask = np.array([[1, 2, 1], [2, 4, 2], [1, 2, 1]])
    new = mean_filter(noise1, mask, calc_arithmetic_mean, 3)
    cv2.imwrite('./res/exp2/res_test2_1_weight_3.jpg', new)
    pprint("使用3*3的加权均值滤波： " + str(mse_loss(ori, new)))

    # 核为3*3的几何均值滤波
    new = mean_filter(noise1, mask, calc_geometric_mean, 3)
    cv2.imwrite('./res/exp2/res_test2_1_mul_3.jpg', new)
    pprint("核为3*3的几何均值滤波： " + str(mse_loss(ori, new)))

    # 核为5*5的几何均值滤波
    new = mean_filter(noise1, mask, calc_geometric_mean, 5)
    cv2.imwrite('./res/exp2/res_test2_1_mul_5.jpg', new)
    pprint("核为5*5的几何均值滤波： " + str(mse_loss(ori, new)))

    # 3*3的谐波均值滤波
    new = mean_filter(noise1, mask, calc_harmonic_mean, 3)
    cv2.imwrite('./res/exp2/res_test2_1_har_3.jpg', new)
    pprint("3*3的谐波均值滤波： " + str(mse_loss(ori, new)))

    # Q为1.5，核大小为3的反谐波均值滤波
    new = mean_filter(noise1, 1.5, calc_contra_harmonic_mean, 3)
    cv2.imwrite('./res/exp2/res_test2_1_char_3_1.5.jpg', new)
    pprint("Q为1.5，核大小为3的反谐波均值滤波： " + str(mse_loss(ori, new)))

    # Q为-1.5，核大小为3的反谐波均值滤波
    new = mean_filter(noise1, -1.5, calc_contra_harmonic_mean, 3)
    cv2.imwrite('./res/exp2/res_test2_1_char_3_-1.5.jpg', new)
    pprint("Q为-1.5，核大小为3的反谐波均值滤波： " + str(mse_loss(ori, new)))

    # 3*3的中值滤波
    new = mean_filter(noise1, mask, calc_median, 3)
    cv2.imwrite('./res/exp2/res_test2_1_median_3.jpg', new)
    pprint("3*3的中值滤波： " + str(mse_loss(ori, new)))

    # 5*5的中值滤波
    new = mean_filter(noise1, mask, calc_median, 5)
    cv2.imwrite('./res/exp2/res_test2_1_median_5.jpg', new)
    pprint("5*5的中值滤波： " + str(mse_loss(ori, new)))

    ori = cv2.imread('./test/test2_2_org.jpg')
    ori = ori.astype('float64')

    noise = cv2.imread('./test/test2_2.jpg')
    noise = noise.astype('float64')

    new = mean_filter(noise, None, calc_median, 3)
    cv2.imwrite('./res/exp2/res_test2_2_median_3.jpg', new)
    pprint('3*3的中值滤波：' + str(mse_loss(ori, new)))
